# 取自52pojie https://www.52pojie.cn/forum.php?mod=viewthread&tid=1977599&extra=page%3D4%26filter%3Dtypeid%26typeid%3D385
import os
import time
import winreg
from selenium import webdriver
from webdriver_manager.chrome import ChromeDriverManager
from webdriver_manager.firefox import GeckoDriverManager
from webdriver_manager.microsoft import IEDriverManager
from webdriver_manager.microsoft import EdgeChromiumDriverManager

BROWSER_PATH = {
    "chrome": r"SOFTWARE\Google\Chrome\BLBeacon",
    "firefox": r"SOFTWARE\mozilla.org\Mozilla",
    "ie": r"SOFTWARE\Microsoft\Internet Explorer",
    "edge": r"SOFTWARE\Microsoft\Edge\BLBeacon",
}

DRIVER_DOWNLOAD = {
    "chrome": (ChromeDriverManager, "chromedriver.exe"),
    "firefox": (GeckoDriverManager, "geckodriver.exe"),
    "ie": (IEDriverManager, "IEDriverServer.exe"),
    "edge": (EdgeChromiumDriverManager, "msedgedriver.exe"),
}

DRIVER_CLASS = {
    "chrome": webdriver.Chrome,
    "firefox": webdriver.Firefox,
    "ie": webdriver.Ie,
    "edge": webdriver.Edge,
}


def get_default_browser():
    key = r"SOFTWARE\Microsoft\Windows\Shell\Associations\UrlAssociations\http\UserChoice"
    try:
        with winreg.OpenKey(winreg.HKEY_CURRENT_USER, key) as reg_key:
            value, _ = winreg.QueryValueEx(reg_key, "Progid")
            for browser in BROWSER_PATH.keys():
                if browser in value.lower():
                    return browser
    except Exception as e:
        print(f"Error retrieving default browser: {e}")
    return None


def get_browser_version(default_browser):
    key = BROWSER_PATH.get(default_browser)
    if key is None:
        return None
    try:
        with winreg.OpenKey(winreg.HKEY_CURRENT_USER, key) as reg_key:
            if default_browser == "firefox":
                value, _ = winreg.QueryValueEx(reg_key, "CurrentVersion")
            else:
                value, _ = winreg.QueryValueEx(reg_key, "version")
            return value
    except Exception as e:
        print(f"Error retrieving browser version for {default_browser}: {e}")
    return None


def check_and_download_driver(browser):
    manager_class, filename = DRIVER_DOWNLOAD.get(browser, (None, None))
    if manager_class is None:
        print(f"Driver for {browser} not supported.")
        return

    # 下载到根目录的drivers文件夹下

    driver_directory = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'drivers', browser)

    if not os.path.exists(driver_directory):
        os.makedirs(driver_directory)

    driver_path = os.path.join(driver_directory, filename)
    if os.path.exists(driver_path):
        print("Driver file already exists at:", driver_path)
        return

    print("Driver file does not exist, start downloading...")
    downloaded_path = manager_class().install()
    print("Driver downloaded successfully at:", downloaded_path)


def open_browser(default_browser):
    driver_class = DRIVER_CLASS.get(default_browser)
    if driver_class is None:
        print(f"No driver class for {default_browser}")
        return

    # driver = driver_class()
    # driver.get("https://www.baidu.com")
    # while True:
    #     time.sleep(1)


if __name__ == "__main__":
    default_browser = get_default_browser()
    print(f"Default Browser: {default_browser}")
    default_browser_version = get_browser_version(default_browser)
    print(f"Default Browser Version: {default_browser_version}")
    check_and_download_driver(default_browser)
    open_browser(default_browser)
    print("Current working directory:", os.getcwd())